/**
 * @file main.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG// 区分CPU调试模式与NPU运行模式
#include "acl/acl.h"
extern void lp_norm_v2_do(uint32_t blockDim, void *stream,
               uint8_t *input, uint8_t *output,  uint8_t *workspace, uint8_t *syncGm, 
               float  p, int32_t axis);
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void lp_norm_v2(GM_ADDR input, GM_ADDR output, 
                                           GM_ADDR workspace, GM_ADDR syncGm, float p, int32_t axis);
#endif

const int32_t DEFAULT_SYNCALL_NEED_SIZE = 8;

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;               // 核数配置（与行数匹配）
    const int32_t rows = 8;              // 矩阵行数
    const int32_t cols = 2048;           // 矩阵列数
    size_t tensorByteSize = rows * cols * sizeof(float);  // 张量总字节数
    float p = 2.0f;                      // 范数阶数（L2范数）
    int32_t axis = -1;                   // 归约轴（-1表示全局范数）
    // size_t workspaceSize = (blockDim + 1) * sizeof(float); // 工作空间大小（存储每个核的局部和）
    size_t workspaceSize = blockDim * 16 * sizeof(float);
    size_t syncWorkspaceSize = (blockDim + 0) * DEFAULT_SYNCALL_NEED_SIZE * sizeof(int32_t);
// CPU模式调试逻辑
#ifdef ASCENDC_CPU_DEBUG
    // 分配全局内存（模拟设备内存）
    uint8_t *input = (uint8_t *)AscendC::GmAlloc(tensorByteSize);
    uint8_t *output = (uint8_t *)AscendC::GmAlloc(tensorByteSize);
    uint8_t *workspace = (uint8_t *)AscendC::GmAlloc(workspaceSize);
    uint8_t *syncGm = (uint8_t *)AscendC::GmAlloc(syncWorkspaceSize); 

    // 读取输入数据（8×2048矩阵）
    ReadFile("./input/input_x.bin", tensorByteSize, input, tensorByteSize);

    // 设置内核运行模式并调用pnorm算子
    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    
    // 直接调用核函数（传递工作空间参数）
    p_norm_v2(input, output, workspace, syncGm, p, axis);

    // 写入输出结果（归一化后的矩阵）
    WriteFile("./output/output_z.bin", output, tensorByteSize);

    // 释放内存
    AscendC::GmFree((void *)input);
    AscendC::GmFree((void *)output);
    AscendC::GmFree((void *)workspace);
    AscendC::GmFree((void *)syncGm);
#else// NPU模式运行逻辑
    // 初始化ACL环境
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));
    
    // 定义主机和设备内存指针（新增sync相关指针）
    uint8_t *input_host, *output_host;
    uint8_t *input_device, *output_device, *workspace_device, *sync_device;

    // 分配主机内存
    CHECK_ACL(aclrtMallocHost((void **)(&input_host), tensorByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&output_host), tensorByteSize));
    
    // 分配设备内存
    CHECK_ACL(aclrtMalloc((void **)&input_device, tensorByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&output_device, tensorByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&workspace_device, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&sync_device, syncWorkspaceSize, ACL_MEM_MALLOC_HUGE_FIRST));  


    // 读取输入文件（8×2048浮点数矩阵）
    ReadFile("./input/input_x.bin", tensorByteSize, input_host, tensorByteSize);

    // 将输入数据从主机拷贝到设备
    CHECK_ACL(aclrtMemcpy(input_device, tensorByteSize, input_host, tensorByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    // 调用pnorm算子（设备端执行），传递工作空间参数
    lp_norm_v2_do(blockDim, stream, input_device, output_device, workspace_device, sync_device, p, axis);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    // 将结果从设备拷贝回主机
    CHECK_ACL(aclrtMemcpy(output_host, tensorByteSize, output_device, tensorByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    
    // 写入输出文件（归一化后的矩阵）
    WriteFile("./output/output_z.bin", output_host, tensorByteSize);

    // 释放设备内存（包含同步缓冲区）
    CHECK_ACL(aclrtFree(input_device));
    CHECK_ACL(aclrtFree(output_device));
    CHECK_ACL(aclrtFree(workspace_device));
    CHECK_ACL(aclrtFree(sync_device));  // 新增：释放同步缓冲区设备内存
    
    // 释放主机内存
    CHECK_ACL(aclrtFreeHost(input_host));
    CHECK_ACL(aclrtFreeHost(output_host));

    // 清理ACL资源
    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}